package org.apache.flink.endpoint;

import io.netty.util.internal.StringUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.ResourceManager;
import org.apache.flink.gateway.JobManagerGateway;
import org.apache.flink.gateway.ResourceManagerGateway;
import org.apache.flink.runtime.rpc.RpcEndpoint;
import org.apache.flink.runtime.rpc.RpcService;
import org.apache.flink.shaded.curator5.org.apache.curator.framework.CuratorFramework;
import org.apache.flink.shaded.curator5.org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.flink.shaded.curator5.org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.flink.shaded.curator5.org.apache.curator.framework.recipes.cache.NodeCache;
import org.apache.flink.shaded.curator5.org.apache.curator.framework.recipes.cache.NodeCacheListener;
import org.apache.flink.shaded.curator5.org.apache.curator.retry.ExponentialBackoffRetry;

import java.util.Timer;
import java.util.TimerTask;

public class JobManagerRpcEndpoint extends RpcEndpoint implements JobManagerGateway {


    private Timer heatbeatTimer;
    private CuratorFramework client;
    private String address;

    public JobManagerRpcEndpoint(RpcService rpcService, String endpointId) throws Exception {
        super(rpcService, endpointId);
        this.heatbeatTimer = new Timer("heatbeat");

        client = CuratorFrameworkFactory.builder()
                .connectString("172.30.18.45:2181")
                .retryPolicy(new ExponentialBackoffRetry(3, 4))
                .namespace("flink")
                .build();
        client.start();


        NodeCache nodeCache = new NodeCache(client, ResourceManager.RM_LEADER_ADDRESS, false);
        nodeCache.getListenable().addListener(new NodeCacheListener() {

            @Override
            public void nodeChanged() throws Exception {
                ChildData childData = nodeCache.getCurrentData();
                address = new String(childData.getData());
            }
        });
        nodeCache.start();

    }

    @Override
    protected void onStart() throws Exception {
        System.out.println("JobManagerRpcEndpoint onStart");

        while (StringUtils.isEmpty(address)) {
            System.out.println("waiting for rm leader");
            Thread.sleep(1000);
        }

        System.out.println("success get rm address: " + address);
        ResourceManagerGateway resourceManagerGateway = getRpcService().connect(address, ResourceManagerGateway.class).get();
        resourceManagerGateway.registerFromJobmanager(getAddress());

        heatbeatTimer.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                ResourceManagerGateway resourceManagerGateway = null;
                try {
                    resourceManagerGateway = getRpcService().connect(address, ResourceManagerGateway.class).get();
                    resourceManagerGateway.heatBeatFromJobmanager(getAddress());
                } catch (Exception e) {
                    System.out.println("error get rm: " + e.getMessage());
                }

            }
        }, 2000, 2000);

    }

    @Override
    public void heatBeatFromResource() {
        System.out.println("JobManagerRpcEndpoint heatBeatFromResource");
    }
}
